package org.finesys.chat.core.base.client;

import org.finesys.ai.core.DeepSeekClient;
import org.finesys.ai.core.StreamingResponseHandling;
import org.finesys.ai.core.SyncOrAsyncOrStreaming;
import org.finesys.ai.core.chat.ChatCompletionRequest;
import org.finesys.ai.core.chat.ChatCompletionResponse;
import org.finesys.ai.core.completion.CompletionRequest;
import org.finesys.ai.core.completion.CompletionResponse;
import org.finesys.ai.core.embedding.EmbeddingRequest;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.finesys.ai.core.image.GenerateImagesRequest;
import org.finesys.ai.core.image.GenerateImagesResponse;
import org.finesys.chat.core.base.chat.ChatRequest;
import org.finesys.chat.core.base.factory.ChatModel;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.Objects;

public class DeepSeekChatClient extends ChatClientAbstract {
    /**
     * deepSeekClient对象
     */
    private DeepSeekClient deepSeekClient;


    public DeepSeekChatClient(ChatModel chatModel) {
        super(chatModel);
        deepSeekClient = build(chatModel);
    }

    /**
     * 根据提供的聊天模型构建DeepSeekClient实例。
     *
     * @param chatModel 聊天模型，包含构建DeepSeekClient所需的配置信息。
     * @return 构建好的DeepSeekClient实例。
     */
    public DeepSeekClient build(ChatModel chatModel) {
        DeepSeekClient.Builder builder = DeepSeekClient.builder()
                .baseUrl(chatModel.getBaseUrl()).model(chatModel.getModel())
                .openAiApiKey(chatModel.getApiKey()).logRequests(chatModel.isLogRequests())
                .logResponses(chatModel.isLogResponses());
        if (Objects.nonNull(chatModel.getProxy())) {
            builder.proxy(chatModel.getProxy());
        }
        if (Objects.nonNull(chatModel.getCallTimeout())) {
            builder.callTimeout(Duration.ofSeconds(chatModel.getCallTimeout()));
        }
        if (Objects.nonNull(chatModel.getConnectTimeout())) {
            builder.connectTimeout(Duration.ofSeconds(chatModel.getConnectTimeout()));
        }
        if (Objects.nonNull(chatModel.getReadTimeout())) {
            builder.readTimeout(Duration.ofSeconds(chatModel.getReadTimeout()));
        }
        if (Objects.nonNull(chatModel.getWriteTimeout())) {
            builder.writeTimeout(Duration.ofSeconds(chatModel.getWriteTimeout()));
        }
        builder.logLevel(chatModel.getLogLevel());
        //注入提示词
        if (chatModel.isDefaultSystemPrompt()) {
            builder.systemMessage(chatModel.getSystemPrompt());
        }
        builder.searchApiKey(chatModel.getSearchApiKey());
        return builder.build();
    }


    @Override
    public void updateKey(ChatModel chatModel) {
        super.updateKey(chatModel);
        if (deepSeekClient == null) {
            throw new RuntimeException("deepSeekClient is null");
        }
        //关闭旧的客户端连接
        deepSeekClient.shutdown();
        //重新构建新的DeepSeekClient实例。
        deepSeekClient = build(chatModel);
    }

    @Override
    public void register(ChatModel chatModel) {
        super.register(chatModel);
        if (deepSeekClient != null) {
            throw new RuntimeException("deepSeekClient is already initialized. Please use updateKey method to update the key.");
        }
        //重新构建新的DeepSeekClient实例。
        deepSeekClient = build(chatModel);
    }

    @Override
    public void unregister() {
        super.unregister();
        if (deepSeekClient == null) {
            throw new RuntimeException("deepSeekClient is null,can not unregister.");
        }
        //关闭旧的客户端连接
        deepSeekClient.shutdown();
    }

    /**
     * 补全接口
     */

    @Override
    public Flux<CompletionResponse> completion(ChatRequest request) {
        CompletionRequest completionRequest = buildCompletionRequest(request);
        if (request.getStream()) {
            return Flux.create(sink -> {
                SyncOrAsyncOrStreaming<CompletionResponse> completionResponse = deepSeekClient.completion(completionRequest);
                //推送每个分片的响应
                StreamingResponseHandling streamingResponseHandling = completionResponse.onPartialResponse(responseContent -> {
                    if (!sink.isCancelled()) {
                        sink.next(responseContent);
                    }
                });
                // 流结束信号
                streamingResponseHandling.onComplete(() -> {
                    if (!sink.isCancelled()) {
                        sink.complete();
                    }
                });
                // 错误处理
                streamingResponseHandling.onError(sink::error);
            });
        }
        //非流式返回
        CompletionResponse completionResponse1 = deepSeekClient.completion(completionRequest).execute();
        return Flux.just(completionResponse1);
    }

    @Override
    public Flux<ChatCompletionResponse> chatCompletion(ChatRequest request) {
        ChatCompletionRequest completionRequest = buildChatCompletionRequest(request);
        return deepSeekClient.chatFluxCompletion(completionRequest);
    }

    @Override
    public Flux<GenerateImagesResponse> imageGeneration(ChatRequest request) {
        GenerateImagesRequest completionRequest = buildImagesCompletionRequest(request);
        //非流式返回
        GenerateImagesResponse imagesResponse = deepSeekClient.imagesGeneration(completionRequest).execute();
        return Flux.just(imagesResponse);
    }

    @Override
    public Flux<EmbeddingResponse> embedding(ChatRequest request) {
        EmbeddingRequest completionRequest = buildEmbeddingRequest(request);
        //非流式返回
        EmbeddingResponse embeddingResponse = deepSeekClient.embedding(completionRequest).execute();
        return Flux.just(embeddingResponse);
    }
}

